package com.lagou.rpc.consumer.client.lb;

import com.alibaba.fastjson.JSON;
import com.lagou.rpc.consumer.client.RpcClient;
import com.lagou.rpc.pojo.RequestStat;
import org.apache.commons.lang3.RandomUtils;
import org.apache.commons.lang3.concurrent.ConcurrentUtils;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.cache.CuratorCache;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * @author xingyi
 */
@Component
public class ClientLoadBalancer {

    private static final Logger logger = LoggerFactory.getLogger(ClientLoadBalancer.class);

    private SortedMap<String, RpcClient> clients = new TreeMap<>();

    private ConcurrentHashMap<String, RequestStat> serverStats = new ConcurrentHashMap<>();

    @Autowired
    private CuratorFramework curatorFramework;

    public void addClient(RpcClient rpcClient) {
        if (clients.putIfAbsent(rpcClient.id(), rpcClient) == null) {
            logger.info("add new client:{}", rpcClient.id());

        }
    }

    public boolean contains(String id) {
        return clients.containsKey(id);
    }

    public void removeClient(String id) {
        if (clients.remove(id) != null) {
            logger.info("remove client:{}", id);
        }
    }

    public RpcClient pick() {
        List<RequestStat> minStats = new ArrayList<>();
        List<RequestStat> sortedStats = serverStats.values()
                .stream()
                .sorted(Comparator.comparing(RequestStat::getConsume))
                .collect(Collectors.toList());
        logger.info("sorted stats:{}", sortedStats);
        RequestStat minStat = sortedStats.get(0);
        for (RequestStat requestStat : sortedStats) {
            if (minStat.getConsume() == requestStat.getConsume()) {
                minStats.add(requestStat);
            } else {
                break;
            }
        }
        final RpcClient rpcClient = clients.get(minStats.get(RandomUtils.nextInt(0, minStats.size())).getServer());
        rpcClient.initClient();
        return rpcClient;
    }

    public void onUpdateStat(String server, byte[] data) {
        RequestStat requestStat = new RequestStat(server, 0, 0);
        if (data != null) {
            requestStat = JSON.parseObject(new String(data, StandardCharsets.UTF_8), RequestStat.class);
        }
        RequestStat stat = ConcurrentUtils.putIfAbsent(serverStats, server, requestStat);
        stat.setConsume(requestStat.getConsume());
        stat.setTimestamp(requestStat.getTimestamp());
    }

    public void reportStat(String server, long consume, long timestamp) throws Exception {
        curatorFramework.setData()
                .forPath(
                        "/rpcserver/" + server,
                        JSON.toJSONString(new RequestStat(server, consume, timestamp)).getBytes(StandardCharsets.UTF_8)
                );
    }

    public RequestStat getStat(String server) {
        return serverStats.get(server);
    }
}
